// Header files
#include <string>
#include <vector>
#include <fstream>
#include <unistd.h>
#include "EFI.h"
#include "UEFI.h"
#include <iostream>
#include <fstream>
#include <map>
#include <sstream>
#include <dirent.h>
#include <stdint.h>
#include "Var.h"
#include <sys/types.h>
#include <sys/stat.h>
#include "ConfigPro.h"
using namespace std;
using namespace config;

// Error codes
enum {IFR_EXTRACTED, FILE_NOT_FOUND, UNKNOWN_PROTOCOL};

// Global definitions
enum type {EFI, UEFI, UNKNOWN};

// Global variables
string fileDir, outputDir, fileLocation, mode, imagePath;
string buffer;
type protocol;
vector<string> strings;
string varGuid = "FFF12B8D-7696-4C8B-A985-2747075B4F50";
string NVARGuid = "CEF5B9A3-476D-497F-9FDC-E98143E0422C";
std::map<int, vector<EFI_VARIABLE_PACK>> VariableMap;
vector<EFI_VARIABLE_PACK>              variablePackages;
vector<string> varFile;
VarConfig Vartest;


std::ostream& operator<<(std::ostream& out, const type value)
{
    map<type, string> strs;

    strs[EFI] = "EFI";
    strs[UEFI] = "UEFI";
    strs[UNKNOWN] = "UNKNOWN";

    return out << strs[value];
}

void getFileNames(string path,vector<string>& filenames)
{
    DIR *pDir;
    struct dirent* ptr;
    if(!(pDir = opendir(path.c_str()))){
        cout << "ERROR: opendir" + path << endl;
        return;
    }
    while((ptr = readdir(pDir))!=0) {
        if (strcmp(ptr->d_name, ".") != 0 && strcmp(ptr->d_name, "..") != 0)
            //filenames.push_back(path + "/" + ptr->d_name);
            filenames.push_back(ptr->d_name);
    }
    closedir(pDir);
}

int CreateDir(const char *sPathName)  
{  
  char DirName[256];  
  strcpy(DirName, sPathName);  
  int i,len = strlen(DirName);
  if (DirName[len -1] != '/')
  {
    DirName[len++] = '/';
    DirName[len] = 0;
  }
  
  for(i=1; i<len; i++)  
  {  
    if(DirName[i]=='/') {  
      DirName[i] = 0; 
      if(access(DirName, F_OK)!=0) { 
#ifdef __WIN32__
		if(mkdir(DirName)==-1) {
#else	  
        if(mkdir(DirName, S_IRUSR | S_IWUSR | S_IXUSR | S_IRWXG | S_IRWXO)==-1) {  
#endif	
          cout << "mkdir   error"  << endl;
          return -1;   
        }  
      }  
      DirName[i] = '/';  
    }  
  }  
  return 0;  
} 

bool dirExists(const string& dirName_in)
{
  bool ret = 0;
  if (0 != access(dirName_in.c_str(), F_OK)) {
    // if this folder not exist, create a new one.
    //ret = mkdir(dirName_in.c_str(),755);   // ���� 0 ��ʾ�����ɹ���-1 ��ʾʧ��
    ret = CreateDir(dirName_in.c_str());
    //���� ::_mkdir  ::_access Ҳ�У���֪��ʲô��˼
  }
  return ret;
}

bool fileExists(const string &file) {

	// Open file
	ifstream fin(file.c_str());
	// Return if first character doesn't equal EOF
	return fin.peek() != EOF;
}

void readFile(const string &file, string &buffer) {

	// Initialize variables
	ifstream fin(file.c_str(), ios::binary);

	// Clear buffer
	buffer.clear();

	// Read in file
	while(fin.peek() != EOF)
		buffer.push_back(fin.get());

	// Close file
	fin.close();
}

type getType(const string &buffer) {

	type currtype = UNKNOWN;


	// Go through buffer
	for(unsigned int i = 0; i < buffer.size() - 9; i++) {

		//cout << "i  " << i << endl;
		// Check if an EFI string package was found
		if((buffer[i] != '\x00' || buffer[i + 1] != '\x00' || buffer[i + 2] != '\x00' || buffer[i + 3] != '\x00') && buffer[i + 4] == '\x02' && buffer[i + 5] == '\x00' && (buffer[i + 6] != '\x00' || buffer[i + 7] != '\x00' || buffer[i + 8] != '\x00' || buffer[i + 9] != '\x00') && 
			(i + static_cast<unsigned char>(buffer[i + 6]) + (static_cast<unsigned char>(buffer[i + 7]) << 8) + (static_cast<unsigned char>(buffer[i + 8]) << 16) + (static_cast<unsigned char>(buffer[i + 9]) << 24) + 4 < buffer.size()) && 
			buffer[i + static_cast<unsigned char>(buffer[i + 6]) + (static_cast<unsigned char>(buffer[i + 7]) << 8) + (static_cast<unsigned char>(buffer[i + 8]) << 16) + (static_cast<unsigned char>(buffer[i + 9]) << 24) + 4] >= 'a' && 
			buffer[i + static_cast<unsigned char>(buffer[i + 6]) + (static_cast<unsigned char>(buffer[i + 7]) << 8) + (static_cast<unsigned char>(buffer[i + 8]) << 16) + (static_cast<unsigned char>(buffer[i + 9]) << 24) + 4] <= 'z' && 
			(i + static_cast<unsigned char>(buffer[i]) + (static_cast<unsigned char>(buffer[i + 1]) << 8) + (static_cast<unsigned char>(buffer[i + 2]) << 16) + (static_cast<unsigned char>(buffer[i + 3]) << 24) + 5) < buffer.size() && 
			buffer[i + static_cast<unsigned char>(buffer[i]) + (static_cast<unsigned char>(buffer[i + 1]) << 8) + (static_cast<unsigned char>(buffer[i + 2]) << 16) + (static_cast<unsigned char>(buffer[i + 3]) << 24) + 4] == '\x02' && 
			buffer[i + static_cast<unsigned char>(buffer[i]) + (static_cast<unsigned char>(buffer[i + 1]) << 8) + (static_cast<unsigned char>(buffer[i + 2]) << 16) + (static_cast<unsigned char>(buffer[i + 3]) << 24) + 5] == '\x00') {
			currtype = EFI;
		} else if(((i + 48) < buffer.size()) && (buffer[i] != '\x00' || buffer[i + 1] != '\x00' || buffer[i + 2] != '\x00') && buffer[i + 3] == '\x04' && buffer[i + 4] == '\x34' && buffer[i + 44] == '\x01' && buffer[i + 45] == '\x00' && buffer[i + 48] == '\x2D' && 
			(i + static_cast<unsigned char>(buffer[i]) + (static_cast<unsigned char>(buffer[i + 1]) << 8) + (static_cast<unsigned char>(buffer[i + 2]) << 16)) < buffer.size() && 
			buffer[i + static_cast<unsigned char>(buffer[i]) + (static_cast<unsigned char>(buffer[i + 1]) << 8) + (static_cast<unsigned char>(buffer[i + 2]) << 16) - 1] == '\x00' &&
			buffer[i + static_cast<unsigned char>(buffer[i]) + (static_cast<unsigned char>(buffer[i + 1]) << 8) + (static_cast<unsigned char>(buffer[i + 2]) << 16) - 2] == '\x00') {
			// Prefer UEFI over EFI
			return UEFI;
		}
	}
	
	return currtype;
}

void defaultVariableInfo(string inPath, vector<string> inVarFile, map<int, vector<EFI_VARIABLE_PACK>> &VarMap) {
  uint64_t offset = 0;
    // Check that input file exists
  ifstream in(inPath.c_str());
  streampos ps;
  uint64_t FileSize = 0;
  char *buf = NULL;
  if(!in.good()){
    cout << "ERROR : "<< inPath <<" not found!!! " << endl;
  }
  in.seekg(0, ios::end);
  ps = in.tellg();
  FileSize = (uint64_t)ps;
  in.seekg(0, std::ios::beg);
  buf = (char *)calloc(1, FileSize+1);
  in.read(buf, FileSize);
  #ifdef DEBUG
  cout << inPath << " FileSize : 0x" << FileSize << endl;
  #endif // DEBUG

  for (uint32_t u32VarIndex = 0; u32VarIndex < VarMap.size(); u32VarIndex++)
  {
    string file = inVarFile[u32VarIndex];
    string::size_type point = file.rfind("tree.bin");
    file.replace(point,8,"info.txt");
    #ifdef DEBUG
    cout <<file << endl;
    #endif // DEBUG
    
    ifstream varfile("./var/" + file);
    string line;

    if(varfile) {// �и��ļ�
      while (getline (varfile, line)) {// line�в�����ÿ�еĻ��з�
          //cout << line << endl;
          point = line.find("Base:");
          if (point != line.npos) {
            line.replace(point,6,"");
            #ifdef DEBUG
            cout << "line : " << line << endl;
            cout << "point = " << point << endl;
            #endif // DEBUG
            offset = std::stoul(line, 0, 16);
            #if 0
            cout << "Variables offset = 0x" << offset << endl;
            #endif // DEBUG
            break; 
          }
      }
    } else {
        cout <<"ERR: no varfile file" << endl;
    }
    #ifdef DEBUG
    cout << "variablePackages[0].variable.offset = " << variablePackages[0].offset << endl;
    cout << "variablePackages[" << variablePackages.size() - 1 << "].offset = " << variablePackages[variablePackages.size() - 1].offset << endl;
    cout << "variablePackages[" << variablePackages.size() - 1 << "].offset = " << variablePackages[variablePackages.size() - 1].offset << endl;
    #endif // DEBUG
  
    memset((buf + offset + VarMap[u32VarIndex][0].offset),
            0xFF, VarMap[u32VarIndex][VarMap[u32VarIndex].size() - 1].offset + \
            VarMap[u32VarIndex][VarMap[u32VarIndex].size() - 1].size - \
            VarMap[u32VarIndex][0].offset);
  }
  

  ofstream outfile("./out.bin", ios::out | ios::binary);
  outfile.write(buf,FileSize);
  outfile.close();
  in.close();
	free(buf);
}

void DelVariableInfo(string inPath, vector<string> inVarFile, map<int, vector<EFI_VARIABLE_PACK>> &VarMap) {
  uint64_t offset = 0;
    // Check that input file exists
  ifstream in(inPath.c_str());
  streampos ps;
  uint64_t FileSize = 0;
  char *buf = NULL;
  if(!in.good()) {
    cout << "ERROR : "<< inPath <<" not found!!! " << endl;
  }  
  in.seekg(0, ios::end);
  ps = in.tellg();
  FileSize = (uint64_t)ps;
  in.seekg(0, std::ios::beg);
  buf = (char *)calloc(1, FileSize+1);
  in.read(buf, FileSize);
  #ifdef DEBUG
  cout << inPath << " FileSize : 0x" << FileSize << endl;
  #endif // DEBUG

  for (uint32_t u32VarIndex = 0; u32VarIndex < VarMap.size(); u32VarIndex++)
  {
    string file = inVarFile[u32VarIndex];
    string::size_type point = file.rfind("tree.bin");
    file.replace(point,8,"info.txt");
    #ifdef DEBUG
    cout <<file << endl;
    #endif // DEBUG
    
    ifstream varfile("./var/" + file);
    string line;

    if(varfile) {// �и��ļ�
      while (getline (varfile, line)) {// line�в�����ÿ�еĻ��з�
          //cout << line << endl;
          point = line.find("Base: ");
          if (point != line.npos) {
            line.replace(point,6,"");
            #ifdef DEBUG
            cout << "line : " << line << endl;
            cout << "point = " << point << endl;
            #endif // DEBUG
            offset = std::stoul(line, 0, 16);
          #ifdef DEBUG
            cout << "Variables offset = 0x" << offset << endl;
            #endif // DEBUG
            break; 
          }
      }
    } else {
        cout <<"ERR: no varfile file" << endl;
    }
  
    for (size_t i = 0; i < VarMap[u32VarIndex].size(); i++)
    {
      if (VarMap[u32VarIndex][i].variableHeader.State != 0x3f)
      {
          continue;
      }
      #ifdef DEBUG
      cout <<"variablePackages[i].variable.data.offset = " << VarMap[u32VarIndex][i].variable.data.offset << endl;
      cout <<"variablePackages[i].variable.data.size = " << VarMap[u32VarIndex][i].variable.data.size << endl;
      #endif // DEBUG
      memcpy((buf + offset + VarMap[u32VarIndex][i].offset),
                &VarMap[u32VarIndex][i].data.at(0), VarMap[u32VarIndex][i].size);
    }
  }
  

  ofstream outfile("./out.bin", ios::out | ios::binary);
  outfile.write(buf,FileSize);
  outfile.close();
  in.close();
	free(buf);
}

void updateVariableInfo(string inPath, vector<string> inVarFile, map<int, vector<EFI_VARIABLE_PACK>> &VarMap) {
  uint64_t offset = 0;
    // Check that input file exists
  ifstream in(inPath.c_str());
  streampos ps;
  uint64_t FileSize = 0;
  char *buf = NULL;
  if(!in.good()) {
    cout << "ERROR : "<< inPath <<" not found!!! " << endl;
  }  
  in.seekg(0, ios::end);
  ps = in.tellg();
  FileSize = (uint64_t)ps;
  in.seekg(0, std::ios::beg);
  buf = (char *)calloc(1, FileSize+1);
  in.read(buf, FileSize);
  #ifdef DEBUG
  cout << inPath << " FileSize : 0x" << FileSize << endl;
  #endif // DEBUG

  for (uint32_t u32VarIndex = 0; u32VarIndex < VarMap.size(); u32VarIndex++)
  {
    string file = inVarFile[u32VarIndex];
    string::size_type point = file.rfind("tree.bin");
    file.replace(point,8,"info.txt");
    #ifdef DEBUG
    cout <<file << endl;
    #endif // DEBUG
    
    ifstream varfile("./var/" + file);
    string line;

    if(varfile) {// �и��ļ�
      while (getline (varfile, line)) {// line�в�����ÿ�еĻ��з�
          //cout << line << endl;
          point = line.find("Base: ");
          if (point != line.npos) {
            line.replace(point,6,"");
            #ifdef DEBUG
            cout << "line : " << line << endl;
            cout << "point = " << point << endl;
            #endif // DEBUG
            offset = std::stoul(line, 0, 16);
          #ifdef DEBUG
            cout << "Variables offset = 0x" << offset << endl;
            #endif // DEBUG
            break; 
          }
      }
    } else {
        cout <<"ERR: no varfile file" << endl;
    }
  
    for (size_t i = 0; i < VarMap[u32VarIndex].size(); i++)
    {
      if (VarMap[u32VarIndex][i].variableHeader.State != 0x3f)
      {
          continue;
      }
      #ifdef DEBUG
      cout <<"variablePackages[i].variable.data.offset = " << VarMap[u32VarIndex][i].variable.data.offset << endl;
      cout <<"variablePackages[i].variable.data.size = " << VarMap[u32VarIndex][i].variable.data.size << endl;
      #endif // DEBUG
      memcpy((buf + offset + VarMap[u32VarIndex][i].variable.data.offset),
                &VarMap[u32VarIndex][i].variable.data.data.at(0), VarMap[u32VarIndex][i].variable.data.size);
    }
  }
  

  ofstream outfile("./out.bin", ios::out | ios::binary);
  outfile.write(buf,FileSize);
  outfile.close();
  in.close();
	free(buf);
}

int generateVarDump(string fileLocation,string outputFile)
{
  int errorCode = IFR_EXTRACTED;
   
  string::size_type point = outputFile.rfind("bin");
  outputFile.replace(point,3,"var");
	// Check if file exists
  
	if(fileExists(fileLocation)) {
    vector<EFI_VOLUME_HEADER_PACK>         volumePackages;
    vector<EFI_VARIABLE_STORE_HEADER_PACK> storePackages;
		readFile(fileLocation, buffer);
		varType type = getVarType(buffer,varGuid);
    if (type == VOLUME_HEADER)
    {
      getVarHeaderInfo( type, buffer, volumePackages);
      #ifdef DEBUG
      displayVarHeaderInfo(volumePackages);
      #endif // DEBUG
      
      //
      //ֻ��һ��variablestoreheader���Ȳ���forѭ��
      //
      getVarStoreHeaderInfo( type, buffer, volumePackages, storePackages);
      #ifdef DEBUG
      displayVarStoreHeaderInfo(storePackages);
      #endif // DEBUG
      getVariableInfo( type, buffer, volumePackages, storePackages,variablePackages);
      #ifdef DEBUG
      displayVariableInfo(variablePackages);
      #endif // DEBUG
    } else {
    //������δʵ��
     
     getNVarInfo(buffer, NVARGuid, variablePackages);
   
    }
  }
  return errorCode;
}

int generateIfrDump(string fileLocation,string outputFile)
{
  int errorCode = IFR_EXTRACTED;
  string::size_type point = outputFile.rfind("bin");
  //outputFile.replace(point,3,"txt");
	// Check if file exists
  if(fileLocation.find("Section_PE32_image_") == string::npos)
  {
      return 1;
  }

	if(fileExists(fileLocation)) {
		readFile(fileLocation, buffer);
		protocol = getType(buffer);
#ifdef DEBUG
    cout << " " << endl;
    cout << "Input: " << fileLocation << endl;
    cout << "Output: " << outputFile << endl;
    cout << "Protocol: " << protocol << endl;
    cout << " " << endl;
#endif // DEBUG
		// Check if protocol is unknown
		if(protocol == UNKNOWN) {

			// Set error code
			errorCode = UNKNOWN_PROTOCOL;

			// Display message
			//cout << "Input: " << fileLocation << " error: Unknown protocol" << endl;

			// Break
			return errorCode;
		}

		// Clear strings
		strings.clear();

		// Check if protocol is EFI
		if(protocol == EFI) {

			// Initialize EFI variables
			vector<EFI_IFR_STRING_PACK> stringPackages;
			vector<EFI_IFR_FORM_SET_PACK> formSets;

			// Get EFI string packages
			getEFIStringPackages(stringPackages, buffer);

			// Get EFI strings
			getEFIStrings(stringPackages, strings, buffer);

			// Get EFI form sets
			getEFIFormSets(formSets, buffer, stringPackages, strings);

			// Generate EFI IFR dump
			//generateEFIIFRDump(outputFile, stringPackages, formSets, buffer, strings);
		}

		// Otherwise check if protocol is UEFI
		else if(protocol == UEFI) {

			// Initialize UEFI variables
			vector<UEFI_IFR_STRING_PACK> stringPackages;
			vector<UEFI_IFR_FORM_SET_PACK> formSets;

			// Get UEFI string packages
			getUEFIStringPackages(stringPackages, buffer);
      #ifdef DEBUG
      displayUEFIStringPackages(stringPackages);
      #endif // DEBUG 
			// Get UEFI strings
			getUEFIStrings(stringPackages, strings, buffer);
      #ifdef DEBUG
      displayUEFIStrings(strings);
      #endif // DEBUG
			// Get UEFI form sets
			getUEFIFormSets(formSets, buffer, stringPackages, strings,fileLocation);
      #ifdef DEBUG
      displayUEFIFormSets(formSets);
      #endif // DEBUG
			// Generate UEFI IFR dump
	    Vartest.generateUEFIIFRDump(outputFile, stringPackages, formSets, buffer, strings, VariableMap, varFile);
      //generateUEFIIFRDump(outputFile, stringPackages, formSets, buffer, strings);
    }
	}

	// Otherwise
	else
	{
		// Set error code
		errorCode = FILE_NOT_FOUND;
		cout << "File not found!" << endl;
	}
  return errorCode;
}

int main(int argc, char **argv)
{
  int errorCode = IFR_EXTRACTED;
  string::size_type pos;
  string FirstMode;
  string SingleBoot;
   if ( argc < 5 )
   {
      cout << "ERR: ifrextract error input para number" << endl;
   }
	
	// Reset error code
	fileDir = argv[1];
  string::size_type point = fileDir.rfind("/");
  if (fileDir.size() == point) {
    fileDir.replace(point,1,"");
    cout << fileDir << endl;
  }
  
	outputDir = argv[2];
  point = outputDir.rfind("/");
  if (outputDir.size() == point) {
    outputDir.replace(point,1,"");
    cout << outputDir << endl;
  }

  imagePath = argv[3];
  mode = argv[4];
  if (argc > 5)
    FirstMode = argv[5];
#if 1
  if ( argc > 6 )
  {   
     if(memcmp(argv[6],"--singleboot=",13)==0)
     {
        SingleBoot = argv[6];
        Vartest.InitSingleBootList(SingleBoot);
     }
  }
   
#endif
  vector<string> files;
  getFileNames(fileDir, files);//�����ļ����ļ��е�·�������?
  int size = files.size();
  if (!size) {
    cout << "ERR: ifrextract input_dir is empty" << endl;
  }
  bool dir_ret = dirExists(outputDir);
  if (!dir_ret)
  {
    /* code */
  }
  
  int SuccessNum = 0;
  varFile.clear();
  VariableMap.clear();

  Vartest.ParseCfgFile();

  for (int i = 0; i<size; i++)
  {
    if (pos != files[i].npos)
    {
      #ifdef DEBUG
      cout << varGuid + " : " + files[i]  << endl;
      #endif // DEBUG
      
      errorCode = generateVarDump(fileDir + "/" + files[i],outputDir + "/" + files[i]);
      if(errorCode){
        cout << "ERROR : dump variable volume fail, please check inputDir!!! " << endl;
        return errorCode;
      }
      
      if(variablePackages.size()) {
        varFile.push_back(files[i]);
        VariableMap[VariableMap.size()] = variablePackages;
        variablePackages.clear();
        //cout << "INFO: variablePackages read from: " << varFile << endl;
        //break;
      }
    }
  }

  for (int i = 0; i<size; i++)
  {
    //pos = files[i].find(varGuid);

    //cout << files[i] << endl;
    errorCode = generateIfrDump(fileDir + "/" + files[i],outputDir + "/" + files[i]);
    if(!errorCode){
      SuccessNum ++;
    }
    //cout << errorCode << endl;
  }

  Vartest.BootModeParse(0, VariableMap);
  Vartest.BootModeParse(1, VariableMap);
  Vartest.BootOrderParse(VariableMap);

  if (mode == "export") {
    Vartest.CreatFile(outputDir + "/config", 0);
  } else if(mode == "import") {
	  if (Vartest.ParseConfig(outputDir + "/config", VariableMap))
    {
      return FILE_NOT_FOUND;
    }

    updateVariableInfo( imagePath, varFile, VariableMap);
    //import var
  }  else if(mode == "default") {
    defaultVariableInfo( imagePath, varFile, VariableMap);
    //import var
  }
  else if(mode == "ManuBootDef") {
    if (Vartest.BootDefault(0, VariableMap))
    {
      return FILE_NOT_FOUND;
    }

    DelVariableInfo( imagePath, varFile, VariableMap);
  }
  else if(mode == "PrimaryBootDef") {
    if (Vartest.BootDefault(1, VariableMap))
    {
      return FILE_NOT_FOUND;
    }

    DelVariableInfo( imagePath, varFile, VariableMap);
  }
  else if (mode == "BootEx") {
    Vartest.CreatFile(outputDir + "/config", BOOTEX);
  }
  else if (mode == "ShowCfg") {
    Vartest.CreatFile(outputDir + "/config", SHOW);
  }
  else if (mode == "InitCfg")
  {
    Vartest.CfgFIleInit(varFile);
  }
  else if (mode == "SetManuFirBoot")
  {
    if (Vartest.SetFirstBoot(0, FirstMode, VariableMap))
    {
      return FILE_NOT_FOUND;
    }
    
    updateVariableInfo( imagePath, varFile, VariableMap);
  }
  else if (mode == "SetPrimaryFirBoot")
  {
    if (Vartest.SetFirstBoot(1, FirstMode, VariableMap))
    {
      return FILE_NOT_FOUND;
    }
    updateVariableInfo( imagePath, varFile, VariableMap);
  }
  else if (std::string::npos != mode.find("FixVar"))
  {
    if (Vartest.FixVar(mode, VariableMap))
    {
      return FILE_NOT_FOUND;
    }
    updateVariableInfo( imagePath, varFile, VariableMap);
  }
//	Vartest.ParseConfig(outputDir + "/", variablePackages);
//	displayVariableInfo(variablePackages);
  if (SuccessNum)
  {
    errorCode = IFR_EXTRACTED;
  }   
	// Return
	return errorCode;
}
